-
Notifications
You must be signed in to change notification settings - Fork 2.6k
Cost display updating for Bedrock custom ARNs that are prompt routers #1604
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Cost display updating for Bedrock custom ARNs that are prompt routers #1604
Conversation
|
|
Think we can clean up some of the logger.debugs? And if the tests that rely on them are important, update them to be based on the behavior or data? |
src/__mocks__/jest.setup.ts
Outdated
| jest.mock("../utils/logging", () => ({ | ||
| logger: { | ||
| debug: jest.fn(), | ||
| debug: jest.fn().mockImplementation((message, meta) => { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we revert the changes to this file and mock debug within specific tests if necessary?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yes, done
yes, unneeded remnant Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
agree, sorry old Javascript habits Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
100d0d3 to
e866079
Compare
e866079 to
8d30b6f
Compare
c8ac41e to
3e15c4e
Compare
src/__mocks__/jest.setup.ts
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Think we can revert this one as well?
3e15c4e to
6c74e9b
Compare
mrubens
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good!
Fix AWS Bedrock Provider with Cost Calculations for Prompt Router ARNs
Note - 2nd PR coming in a couple hours to make AWS prices in src/shared/api.ts current. These 2 PRs together will make Roo's cost display more accurate for Bedrock users.
Overview
This PR fixes the AWS Bedrock provider to properly handle and utilize the
invokedModelIdinformation returned by AWS Bedrock Prompt Router responses. The code has been updated to use the token cost information from the invokedModelId, because at request time the inference model, and therefore token costs, are not known for prompt router ARNs. This improvement ensures more accurate cost tracking and model configuration when using AWS Bedrock's Prompt Router feature.Problem
When using AWS Bedrock's Prompt Router feature, the actual model used for inference is not known at request time. The Prompt Router returns the actual model used via the
invokedModelIdfield in the response stream, and the token usage values are located in a different structure than responses from individual models. Previously, our implementation wasn't properly extracting and utilizing this information, which led to:Changes
invokedModelIdfrom AWS Bedrock Prompt Router responsesBenefits
Testing
Added new test files and expanded existing tests:
bedrock-invokedModelId.test.ts: Tests for proper handling of the invokedModelId fieldbedrock-custom-arn.test.ts: Tests for custom ARN validation and region extractionbedrock.test.ts: Added tests for invokedModelId handling in various scenariosTechnical Details
The key technical changes include:
This implementation ensures that when AWS Bedrock's Prompt Router selects a different model than initially requested, our system correctly identifies and uses that model's configuration for accurate cost tracking and proper functionality.
Important
Enhances AWS Bedrock provider to accurately handle cost calculations and model configurations for prompt router ARNs, with improved error handling and comprehensive testing.
AwsBedrockHandlerinbedrock.tsto handleinvokedModelIdfor accurate cost tracking and model configuration.invokedModelIdand updatecostModelConfig.api.tswith latest pricing and new models.supportsComputerUseflag to relevant models.bedrock-invokedModelId.test.tsforinvokedModelIdhandling.bedrock-createMessage.test.tsandbedrock.test.tsfor custom ARN and cross-region inference scenarios.jest.setup.tsfor consistent test logging.This description was created by
for d8df9a5. It will automatically update as commits are pushed.